package org.apache.lorry.zookeeper.cluster;

import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.util.ArrayList;
import java.util.List;

/**
 * 恢复分配的任务
 */
@Slf4j
public class RecoveredAssignments {
    /*
     * Various lists wew need to keep track of.
     */
    List<String> tasks;
    List<String> assignments;
    List<String> status;
    List<String> activeWorkers;
    List<String> assignedWorkers;

    RecoveryCallback cb;

    ZooKeeper zk;

    /**
     * Callback interface. Called once
     * recovery completes or fails.
     *
     */
    public interface RecoveryCallback {
        final static int OK = 0;
        final static int FAILED = -1;

        public void recoveryComplete(int rc, List<String> tasks);
    }

    public RecoveredAssignments(ZooKeeper zk) {
        this.zk = zk;
        this.assignments = new ArrayList<>();
    }


    public void recover(RecoveryCallback cb) {
        this.cb = cb;
        getTasks();
    }

    private void getTasks() {
        zk.getChildren(ZooConstant.TASK_PATH, false, tasksGetChildrenCallback, null);
    }

    AsyncCallback.ChildrenCallback tasksGetChildrenCallback = new AsyncCallback.ChildrenCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, List<String> children) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getTasks();
                    break;
                case OK:
                    tasks = children;
                    getAssignedWorkers();
                    break;
                default:
                    log.error("getChildren failed",  KeeperException.create(KeeperException.Code.get(rc), path));
                    cb.recoveryComplete(RecoveryCallback.FAILED, null); //失败

            }
        }
    };

    //获取已经分配的worker
    private void getAssignedWorkers(){
        zk.getChildren(ZooConstant.ASSIGN_PATH, false, assignedWorkersGetChildrenCallback, null);
    }


    AsyncCallback.ChildrenCallback assignedWorkersGetChildrenCallback = new AsyncCallback.ChildrenCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, List<String> children) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getAssignedWorkers();
                    break;
                case OK:
                    assignedWorkers = children;
                    getWorkers(children);
                    break;

                default:
                    log.error("getChildren failed", KeeperException.create(KeeperException.Code.get(rc), path));
                    cb.recoveryComplete(RecoveryCallback.FAILED, null); //失败
            }
        }
    };

    private void getWorkers(List<String> assignedWorkers) {
        zk.getChildren(ZooConstant.WORKER_PATH, false, workersCallback, assignedWorkers);
    }


    AsyncCallback.ChildrenCallback workersCallback = new AsyncCallback.ChildrenCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, List<String> children) {
            switch(KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getWorkers((List<String>) ctx);
                    break;
                case OK:
                    if (CollectionUtils.isNotEmpty(children)) { //如果为空，没有有效的worker
                        log.info("Empty list of workers, possibly just starting");
                        cb.recoveryComplete(RecoveryCallback.OK, new ArrayList<>());
                        break;
                    }
                    activeWorkers = children;   //获取所有在线的worker
                    for (String s : assignedWorkers) {
                        getWorkerAssignments(ZooConstant.ASSIGN_PATH + "/" + s);
                    }
                    break;
                default:
                    log.error("getChildren failed", KeeperException.create(KeeperException.Code.get(rc), path));
                    cb.recoveryComplete(RecoveryCallback.FAILED, null); //失败
            }
        }
    };


    private void getWorkerAssignments(String assignWorkPath) {//获取分配的任务
        zk.getChildren(assignWorkPath, false, workerAssignmentsCallback, null);
    }

    AsyncCallback.ChildrenCallback workerAssignmentsCallback = new AsyncCallback.ChildrenCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, List<String> children) {
            switch(KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getWorkerAssignments(path);
                    break;
                case OK:
                    String workName = path.substring(path.lastIndexOf("/") + 1);
                    if(activeWorkers.contains(workName)) { //如果在线
                        assignments.addAll(children);   //将任务全部添加到assignments中
                    }else{  //此worker已经下线，它下面的Task需要重新分配
                        for (String task : children) {
                            if (!tasks.contains(task)) {    //任务不在任务列表里面了。
                                tasks.add(task);
                                getDataReassign(path, task);
                            }else{
                                deleteAssignment(path + "/" + task);
                            }
                        }
                        deleteAssignment(path);
                    }

                    assignedWorkers.remove(workName);

                    /*
                     * Once we have checked all assignments,
                     * it is time to check the status of tasks
                     */
                    if(assignedWorkers.isEmpty()) {
                        log.info("Getting statuses for recovery");
                        getTasksStatus();
                    }
                    break;

                case NONODE:
                    log.info("No such znode exists: " + path);
                    break;
                default:
                    log.error("getChildren failed", KeeperException.create(KeeperException.Code.get(rc), path));
                    cb.recoveryComplete(RecoveryCallback.FAILED, null); //失败
            }
        }
    };

    private void getTasksStatus() {
        zk.getChildren(ZooConstant.STATUS_PATH, false, tasksStatusCallback, null);
    }

    AsyncCallback.ChildrenCallback tasksStatusCallback = new AsyncCallback.ChildrenCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, List<String> children) {
            switch(KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getTasksStatus();
                    break;
                case OK:
                    status = children;
                    processAssignments();
                    break;
                default:
                    log.error("getChildren failed", KeeperException.create(KeeperException.Code.get(rc), path));
                    cb.recoveryComplete(RecoveryCallback.FAILED, null); //失败
            }
        }
    };


    private void processAssignments(){
        log.info("Size of tasks: " + tasks.size());
        for(String assignment : assignments){
            deleteAssignment(ZooConstant.TASK_PATH + "/" + assignment);//删除已经分配了的
            tasks.remove(assignment);
        }

        log.info("Size of tasks after assignment filtering: " + tasks.size());

        for(String s: status){
            deleteAssignment(ZooConstant.TASK_PATH + "/" + s);//删除已经分配了的
            tasks.remove(s);
        }
        log.info("Size of tasks after status filtering: " + tasks.size());
        cb.recoveryComplete(RecoveryCallback.OK, tasks);
    }

    //重新进行分配
    void getDataReassign(String path, String task) {
        //获取分配的task的信息
        zk.getData(path + "/" + task, false, getDataReassignCallback, task);
    }

    AsyncCallback.DataCallback getDataReassignCallback = new AsyncCallback.DataCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, byte[] data, Stat stat) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getDataReassign(path, (String) ctx);
                case OK:
                    RecreateTaskCtx recreateTaskCtx = new RecreateTaskCtx(path, (String) ctx, data);
                    recreateTask(recreateTaskCtx);
                    break;
                default:
                    log.error("Something went wrong when getting data ", KeeperException.create(KeeperException.Code.get(rc)));
            }
        }
    };


    void recreateTask(RecreateTaskCtx ctx) {
        zk.create(ZooConstant.TASK_PATH + "/" + ctx.task,
                ctx.data,
                ZooDefs.Ids.OPEN_ACL_UNSAFE,
                CreateMode.PERSISTENT,
                recreateTaskCallback,
                ctx);
    }

    AsyncCallback.StringCallback recreateTaskCallback = new AsyncCallback.StringCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, String name) {
            RecreateTaskCtx taskCtx = (RecreateTaskCtx) ctx;
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    recreateTask(taskCtx);
                    break;
                case OK:
                    deleteAssignment(taskCtx.path);
                    break;
                default:
                    log.error("Something wwnt wrong when recreating task", KeeperException.create(KeeperException.Code.get(rc)));
            }
        }
    };

    void deleteAssignment(String path){
        zk.delete(path, -1, taskDeleteCallback, null);
    }

    AsyncCallback.VoidCallback taskDeleteCallback = new AsyncCallback.VoidCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx) {
            switch(KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    deleteAssignment(path);
                    break;
                case OK:
                    log.info("Task correctly deleted: {}", path);
                    break;
                default:
                    log.error("Failed to delete task data" + KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }
    };



    class RecreateTaskCtx {
        String path;
        String task;
        byte[] data;

        RecreateTaskCtx(String path, String task, byte[] data) {
            this.path = path;
            this.task = task;
            this.data = data;
        }
    }

}